import os
import argparse
from db_utils import InstanceDatabase

# CONFIGURABLE PATHS
# CLUSTER_DIR = "clustering_results"  # directory containing clusterXX.txt
# DB_PATH = "meta.db"                 # path to your SQLite DB
# OUTPUT_DIR = "tracks_output"        # where to create track folders

def main():
    parser = argparse.ArgumentParser(description="Split cluster files by track using DB info.")
    parser.add_argument('--cluster_dir', type=str, default="clustering_results", help='Directory containing clusterXX.txt files')
    parser.add_argument('--db_path', type=str, default="meta.db", help='Path to the SQLite DB')
    parser.add_argument('--output_dir', type=str, default="tracks_output", help='Where to create track folders')
    args = parser.parse_args()

    CLUSTER_DIR = args.cluster_dir
    DB_PATH = args.db_path
    OUTPUT_DIR = args.output_dir

    # Step 1: Connect to DB and get track info
    db = InstanceDatabase(DB_PATH)
    tracks = db.get_available_tracks()

    # Step 2: Read all cluster files
    cluster_files = [f for f in os.listdir(CLUSTER_DIR) if f.startswith("cluster") and f.endswith(".txt")]
    cluster_to_hashids = {}
    for fname in cluster_files:
        with open(os.path.join(CLUSTER_DIR, fname)) as f:
            hashids = set(line.strip() for line in f if line.strip())
            cluster_to_hashids[fname] = hashids

    print(list(cluster_to_hashids.keys()))
    # Step 3: For each track, create folder and write relevant cluster files
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for track in tracks:
        track_hashids = db.get_track_uuids(track)
        track_folder = os.path.join(OUTPUT_DIR, track)
        os.makedirs(track_folder, exist_ok=True)
        for cluster_fname, cluster_hashids in cluster_to_hashids.items():
            relevant_hashids = cluster_hashids & track_hashids
            if relevant_hashids:
                with open(os.path.join(track_folder, cluster_fname), "w") as out_f:
                    for hashid in sorted(relevant_hashids):
                        out_f.write(hashid + "\n")
    db.close()
    print("Done. Check the", OUTPUT_DIR, "directory.")

if __name__ == "__main__":
    main() 